-
Notifications
You must be signed in to change notification settings - Fork 598
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[compiler][flow] Move cast, reshape and bitcast after transfer op #18742
base: main
Are you sure you want to change the base?
Conversation
We got incoming IR of the form ```mlir %cast = tensor.cast %0 : tensor<2xf32> to tensor<?xf32> %2 = flow.tensor.transfer %cast : tensor<?xf32>{%c2} to #hal.device.affinity<@__device_0> %cast_2 = tensor.cast %2 : tensor<?xf32> to tensor<2xf32> ``` We would like to allow for the 2 casts to get folded. We would also like to reduce the dynamism of the transfer op. To operate on a tensor with fewer dynamic dimensions. Signed-off-by: Boian Petkantchin <boian.petkantchin@amd.com>
f4c155a
to
d5794d3
Compare
I've been trying to fix a similar problem with If this is a recurring problem, maybe there is a more systematic way to address this? |
@IanWood1 This seems similar. I am not sure what is more appropriate, to fold the reshape or to move it. |
Folding is best whenever possible - the reshape/cast ops are only there to carry metadata and should be aggressively removed when the metadata they carry is not useful. For example, %0 = flow.tensor.reshape %arg : tensor<1x2xf32> -> tensor<2x1xf32>
%1 = flow.tensor.transfer %0 : tensor<2x1xf32> to "target"
%2 = flow.tensor.reshape %1 : tensor<2x1xf32> -> tensor<1x2xf32> should never exist - the canonical form would be: %2 = flow.tensor.transfer %arg : tensor<1x2xf32> to "target" If propagating the reshapes/casts/etc allows those to fold by having an intermediate state like: %0 = flow.tensor.transfer %arg : tensor<1x2xf32> to "target"
%1 = flow.tensor.reshape %0 : tensor<1x2xf32> -> tensor<2x1xf32>
%2 = flow.tensor.reshape %1 : tensor<2x1xf32> -> tensor<1x2xf32> and then the existing reshape-reshape canonicalizer kicks in that's a good thing. |
|
||
// Unranked shapes are always considered to have more dynamic dimensions than | ||
// ranked. | ||
inline bool shapeHasLessOrEqualDynamicDimensions(ShapedType t1, ShapedType t2) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
YAGNI
on the 3rd or 4th time a 3 line block of code is used it's worth hoisting into a global shared util - but there's a bar for doing this - every shared util pollutes the project and adds a maintenance burden and has to be worth it. A single use in a single location is not worth it. Just inline this as static where this is used.
(this would also need to be static here - you can't put inline functions in header files)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I moved it.
template <typename Op, unsigned homomorphismOpOperandIndex = 0, | ||
unsigned homomorphismOpResultIndex = 0> | ||
static void populateMoveOpAfterTransferPattern(RewritePatternSet &results) { | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
// before the transfer instead. We strive to reduce the dynamism of the | ||
// transfer op. If there will be no strict dynamism improvement, we prefer the | ||
// other op after the transfer. | ||
// TODO: add the analogous move-befor-transfer pattern. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
// TODO: add the analogous move-befor-transfer pattern. | |
// TODO: add the analogous move-before-transfer pattern. |
(a spell checker can help in your IDE)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done. I have one but it ignored the hyphened word.
} // namespace | ||
|
||
void TensorTransferOp::getCanonicalizationPatterns(RewritePatternSet &results, | ||
MLIRContext *context) { | ||
results.insert<ElideRedundantTransfer>(context); | ||
populateMoveOpAfterTransferPattern<tensor::BitcastOp>(results); | ||
populateMoveOpAfterTransferPattern<TensorBitCastOp>(results); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
populateMoveOpAfterTransferPattern<TensorBitCastOp>(results); | |
populateMoveOpAfterTransferPattern<IREE::Flow::TensorBitCastOp>(results); |
use namespaced names where there may be ambiguity to make it clearer to readers
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
populateMoveOpAfterTransferPattern<TensorBitCastOp>(results); | ||
populateMoveOpAfterTransferPattern<tensor::CastOp>(results); | ||
populateMoveOpAfterTransferPattern<tensor::ReshapeOp>(results); | ||
populateMoveOpAfterTransferPattern<TensorReshapeOp>(results); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
populateMoveOpAfterTransferPattern<TensorReshapeOp>(results); | |
populateMoveOpAfterTransferPattern<IREE::Flow::TensorReshapeOp>(results); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
We got incoming IR of the form
We would like to allow for the 2 casts to get folded. We would also like to reduce the dynamism of the transfer op. To operate on a tensor with fewer dynamic dimensions.